/*
 * @(#)SOAPSignatureHandler.java
 * Date 2013-02-24
 * Version 1.0
 * Author Jim Horner
 * Copyright (c)2012
 */


package us.hornerscorners.lamppost.ws;

import java.util.ArrayList;
import java.util.Calendar;
import java.util.Collections;
import java.util.GregorianCalendar;
import java.util.List;
import java.util.Set;

import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.Transform;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureFactory;

import javax.xml.crypto.dsig.dom.DOMSignContext;

import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;

import javax.xml.datatype.DatatypeFactory;
import javax.xml.datatype.XMLGregorianCalendar;

import javax.xml.namespace.QName;

import javax.xml.soap.SOAPBody;
import javax.xml.soap.SOAPElement;
import javax.xml.soap.SOAPEnvelope;
import javax.xml.soap.SOAPHeader;
import javax.xml.soap.SOAPHeaderElement;
import javax.xml.soap.SOAPMessage;
import javax.xml.soap.SOAPPart;

import javax.xml.ws.handler.MessageContext;

import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;

import org.apache.commons.codec.binary.Base64;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 *
 * @author jim
 */
public class SOAPSignatureHandler implements SOAPHandler<SOAPMessageContext> {

    /** Field description */
    private static final String ATTRIBUTENAME_X509TOKEN =
        "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-x509-token-profile-1.0#X509v3";

    /** Field description */
    private static final String ENCODINGTYPE_BASE64 =
        "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-soap-message-security-1.0#Base64Binary";

    /** Field description */
    private static final String NAMESPACEURI_WSSECURITY_WSSE =
        "http://docs.oasisopen.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd";

    /** Field description */
    private static final String NAMESPACEURI_WSSECURITY_WSU =
        "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd";

    /** Field description */
    private static final String NAMESPACEURI_XMLSIGNATURE_DS =
        "http://www.w3.org/2000/09/xmldsig#";

    /** Field description */
    private final KeyPair keypair;

    /** Field description */
    private final Logger logger =
        LoggerFactory.getLogger(SOAPDebugHandler.class);

    /**
     * Constructs ...
     *
     *
     */
    public SOAPSignatureHandler(KeyPair keypair) {

        super();
        this.keypair = keypair;
    }

    /**
     *  Method description
     *
     *
     *  @param context
     */
    @Override
    public void close(MessageContext context) {

        // throw new UnsupportedOperationException("Not supported yet.");
    }

    /**
     * Method description
     *
     *
     * @return
     */
    @Override
    public Set<QName> getHeaders() {

        return null;
    }

    /**
     * Method description
     *
     *
     * @param context
     *
     * @return
     */
    @Override
    public boolean handleFault(SOAPMessageContext context) {

        // throw new UnsupportedOperationException("Not supported yet.");
        return true;
    }

    /**
     * Method description
     *
     *
     * @param smc
     *
     * @return
     */
    @Override
    public boolean handleMessage(SOAPMessageContext smc) {

        // TODO investigate better ways, this code was copied off internet in hurry
        // http://docs.oracle.com/javase/7/docs/technotes/guides/security/xmldsig/GenEnveloping.java
        // http://docs.oracle.com/javase/7/docs/technotes/guides/security/xmldsig/GenEnveloped.java



        Boolean outboundProperty =
            (Boolean) smc.get(MessageContext.MESSAGE_OUTBOUND_PROPERTY);

        SOAPMessage msg = smc.getMessage();

        try {

            if (Boolean.TRUE.equals(outboundProperty)) {

//              Prepare envelope
                SOAPPart soapPart = msg.getSOAPPart();
                SOAPEnvelope envelope = soapPart.getEnvelope();

                envelope.addNamespaceDeclaration("u",
                                                 NAMESPACEURI_WSSECURITY_WSU);

//              Prepare header
                SOAPHeader header = envelope.getHeader();

                if (header == null) {
                    header = envelope.addHeader();
                }

                SOAPHeaderElement securityElement =
                    header.addHeaderElement(
                        new QName(
                            NAMESPACEURI_WSSECURITY_WSSE, "Security", "o"));

                securityElement.setMustUnderstand(true);

//              Prepare body
                SOAPBody body = envelope.getBody();

                body.addAttribute(new QName(NAMESPACEURI_WSSECURITY_WSU, "Id",
                                            "u"), "body");

//              Generate timestamps
                GregorianCalendar calendar = new GregorianCalendar();
                XMLGregorianCalendar xmlNow =
                    DatatypeFactory.newInstance().newXMLGregorianCalendar(
                        calendar);

                calendar.add(Calendar.MINUTE, 5);

                XMLGregorianCalendar xmlLater =
                    DatatypeFactory.newInstance().newXMLGregorianCalendar(
                        calendar);

//              Prepare timestamp element
                SOAPElement timeStampElement =
                    securityElement.addChildElement("Timestamp", "u");

                timeStampElement.addAttribute(
                    new QName(NAMESPACEURI_WSSECURITY_WSU, "Id", "u"),
                    "timestamp");

                SOAPElement createdElement =
                    timeStampElement.addChildElement("Created", "u");

                createdElement.addTextNode(xmlNow.toString());

                SOAPElement expiresElement =
                    timeStampElement.addChildElement("Expires", "u");

                expiresElement.addTextNode(xmlLater.toString());

//              Prepare security token element
                SOAPElement binarySecurityTokenElement =
                    securityElement.addChildElement("BinarySecurityToken", "o");

                binarySecurityTokenElement.addAttribute(
                    new QName(NAMESPACEURI_WSSECURITY_WSU, "Id", "u"), "cert");
                binarySecurityTokenElement.addAttribute(new QName("ValueType"),
                        ATTRIBUTENAME_X509TOKEN);
                binarySecurityTokenElement.addTextNode(
                    new String(
                        Base64.encodeBase64(
                            this.keypair.getPublicCert().getEncoded())));

//              Signature generation
                XMLSignatureFactory signFactory =
                    XMLSignatureFactory.getInstance("DOM");
                C14NMethodParameterSpec spec1 = null;
                CanonicalizationMethod c14nMethod =
                    signFactory.newCanonicalizationMethod(
                        CanonicalizationMethod.EXCLUSIVE, spec1);
                DigestMethod digestMethod =
                    signFactory.newDigestMethod(DigestMethod.SHA1, null);
                SignatureMethod signMethod =
                    signFactory.newSignatureMethod(SignatureMethod.RSA_SHA1,
                                                   null);
                TransformParameterSpec spec2 = null;
                Transform transform =
                    signFactory.newTransform(CanonicalizationMethod.EXCLUSIVE,
                                             spec2);
                List<Transform> transformList =
                    Collections.singletonList(transform);
                List<Reference> referenceList = new ArrayList<>();
                Reference reference1 = signFactory.newReference("#body",
                                           digestMethod, transformList, null,
                                           null);

                referenceList.add(reference1);

                Reference reference2 = signFactory.newReference("#timestamp",
                                           digestMethod, transformList, null,
                                           null);

                referenceList.add(reference2);

                SignedInfo signInfo = signFactory.newSignedInfo(c14nMethod,
                                          signMethod, referenceList);
                DOMSignContext dsc =
                    new DOMSignContext(this.keypair.getPrivateKey(),
                                       securityElement);
                XMLSignature signature = signFactory.newXMLSignature(signInfo,
                                             null);

                signature.sign(dsc);

                // Prepare key info element
                SOAPElement signatureElement =
                    (SOAPElement) securityElement.getLastChild();

                SOAPElement keyInfoElement =
                    signatureElement.addChildElement("KeyInfo");

                SOAPElement securityTokenReferenceElement =
                    keyInfoElement.addChildElement("SecurityTokenReference",
                                                   "o");

                SOAPElement referenceElement =
                    securityTokenReferenceElement.addChildElement("Reference",
                        "o");

                referenceElement.setAttribute("URI", "#cert");

            }

        } catch (Exception ex) {

            logger.debug("Error in the client handler", ex);
        }

        return true;
    }
}
